# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Lenet Tutorial
The sample can be run on CPU, GPU and Ascend 910 AI processor.
"""
import os
from urllib.parse import urlparse
import argparse
import mindspore
import mindspore.dataset as ds
import mindspore.nn as nn
from mindspore import context
from mindspore.train.serialization import load_checkpoint, load_param_into_net
from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor
from mindspore.train import Model
from mindspore.common.initializer import TruncatedNormal
import mindspore.dataset.transforms.vision.c_transforms as CV
import mindspore.dataset.transforms.vision.py_transforms as PV
from mindvision
import mindspore.dataset.transforms.c_transforms as C
import mindspore.dataset.transforms.py_transforms as PT
from mindspore.nn.metrics import Accuracy
from mindspore.common import dtype as mstype
from mindspore.nn.loss import SoftmaxCrossEntropyWithLogits


def create_dataset_from_folder(data_path, batch_size=32, repeat_size=1,
                   num_parallel_workers=1):
    """ create dataset for train or test
    Args:
        data_path: Data path
        batch_size: The number of data records in each group
        repeat_size: The number of replicated data records
        num_parallel_workers: The number of parallel workers
    """
    # define dataset
    mnist_ds = ds.ImageFolderDatasetV2(dataset_dir=data_path, decode=False)
    
    # define operation parameters
    resize_height, resize_width = 32, 32
    nml_mean = [0.1307]
    nml_std = [0.3081]

    transform = [
        PV.Decode(),
        PV.Grayscale(1),
        PV.Resize(size=(resize_height, resize_width)),
        PV.ToTensor(),
        PV.Normalize(nml_mean, nml_std),
        ]
    compose = PV.ComposeOp(transform)

    # apply map operations on images
    mnist_ds = mnist_ds.map(input_columns="label", operations=C.TypeCast(mindspore.int32))
    mnist_ds = mnist_ds.map(input_columns="image", operations=compose())

    # apply DatasetOps
    buffer_size = 10000
    mnist_ds = mnist_ds.shuffle(buffer_size=buffer_size)  # 10000 as in LeNet train script
    mnist_ds = mnist_ds.batch(batch_size, drop_remainder=True)
    mnist_ds = mnist_ds.repeat(repeat_size)

    return mnist_ds

def conv(in_channels,out_channels,kernel_size,stride=1,padding=0):
    """Conv layer weight initial"""
    weight = weight_variable()
    return nn.Conv2d(in_channels,out_channels,kernel_size=kernel_size,
                     stride=stride,padding=padding,weight_init=weight,has_bias=False,pad_mode="valid")
def fc_with_initialize(input_channels,out_channels):
    """FC layer weigh initial"""
    weight = weight_variable()
    bias = weight_variable()
    return nn.Dense(input_channels,out_channels,weight,bias)
def weight_variable():
    """Weight initial."""
    return TruncatedNormal(0.02)
# TODO: design the Net
class Net(nn.Cell):
    """Net network structure."""
    def __init__(self):
        super(Net, self).__init__()
        # TODO: define the operator required
        self.conv1 = conv(1, 6, 5)
        self.conv2 = conv(6, 16, 5)
        self.fc1 = fc_with_initialize(16 * 5 * 5, 120)
        self.fc2 = fc_with_initialize(120, 84)
        self.fc3 = fc_with_initialize(84, 10)
        self.relu = nn.ReLU()
        self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2)
        self.flatten = nn.Flatten()

    def construct(self, x):
        # TODO: use the preceding operators to construct networks
        x = self.conv1(x)
        x = self.relu(x)
        x = self.max_pool2d(x)
        x = self.conv2(x)
        x = self.relu(x)
        x = self.max_pool2d(x)
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.relu(x)
        x = self.fc3(x)
        return x


def train_net(args, model, epoch_size, mnist_path, repeat_size, ckpoint_cb):
    """Define the training method."""
    print("============== Starting Training ==============")
    # load training dataset
    ds_train = create_dataset_from_folder(os.path.join(mnist_path, "train"), 32, repeat_size)
    model.train(epoch_size, ds_train, callbacks=[ckpoint_cb, LossMonitor()], dataset_sink_mode=False)


def test_net(args, network, model, mnist_path):
    """Define the evaluation method."""
    print("============== Starting Testing ==============")
    # load the saved model for evaluation
    #TODO modify the checkpoint file name
    param_dict = load_checkpoint("mnist.ckpt")
    # load parameter to the network
    load_param_into_net(network, param_dict)
    # load testing dataset
    ds_eval = create_dataset_from_folder(os.path.join(mnist_path, "test"))
    acc = model.eval(ds_eval, dataset_sink_mode=False)
    print("============== Accuracy:{} ==============".format(acc))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='MindSpore Mnist Example')
    parser.add_argument('--device_target', type=str, default="CPU", choices=['Ascend', 'GPU', 'CPU'],
                        help='device where the code will be implemented (default: CPU)')
    args = parser.parse_args()
    context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target)
    # learning rate setting
    lr = 0.01
    momentum = 0.9
    epoch_size = 5
    mnist_path = "./mnist"

    # TODO: construct a Model to train the Net
    # define the loss function
    # create the network
    # define the optimizer
    # save the network model and parameters for subsequence fine-tuning
    # group layers into an object with training and evaluation features
    net_loss = SoftmaxCrossEntropyWithLogits(is_grad=False,sparse=True,reduction='mean')
    repeat_size = 1
    network = Net()
    net_opt = nn.Momentum(network.trainable_params(),lr,momentum)
    config_ck = CheckpointConfig(save_checkpoint_steps=1875,keep_checkpoint_max=10)
    ckpoint_cb = ModelCheckpoint(prefix="mnist",config=config_ck)
    model = Model(network,net_loss,net_opt,metrics={"Accuracy":Accuracy()})

    train_net(args, model, epoch_size, mnist_path, repeat_size, ckpoint_cb)
    #test_net(args, network, model, mnist_path)
